import java.util.concurrent.*
import java.util.logging.Logger
import org.semanticweb.owlapi.apibinding.OWLManager
import org.semanticweb.owlapi.model.*
import org.semanticweb.owlapi.reasoner.*
import org.semanticweb.owlapi.profiles.*
import org.semanticweb.owlapi.util.*
import org.mindswap.pellet.KnowledgeBase
import org.mindswap.pellet.expressivity.*
import org.mindswap.pellet.*
import org.semanticweb.owlapi.io.*
import org.semanticweb.elk.owlapi.*

def THREADS = 10

pool = Executors.newFixedThreadPool(THREADS)
defer = { c -> pool.submit(c as Callable) }

def cli = new CliBuilder()
cli.with {
usage: 'foobar'
  h longOpt:'help', 'this information'
  g longOpt:'goa-file', 'input file containing GO annotations', args:1, required:true
  n longOpt:'new-annos', 'input file containing predicted GO annotations',args:1, required:true
  o longOpt:'output', 'output file containing similarity values',args:1, required:true
  //  t longOpt:'threads', 'number of threads', args:1
  //  k longOpt:'stepsize', 'steps before splitting jobs', arg:1
}

def opt = cli.parse(args)
if( !opt ) {
  //  cli.usage()
  return
}
if( opt.h ) {
  cli.usage()
    return
}

def interactions = new File("gene_genetic_interactions_fb_2012_01.tsv")

def orthomap = [:]
interactions.splitEachLine("\t") { line ->
  if (line[0].indexOf('#')==-1 && line.size()>1) {
    def p1 = line[1]
    def p2 = line[3]
    if (orthomap[p1] == null) {
      orthomap[p1] = new TreeSet()
    }
    orthomap[p1].add(p2)
  }
}

OWLOntologyManager manager = OWLManager.createOWLOntologyManager()

OWLDataFactory fac = manager.getOWLDataFactory()
def factory = fac

OWLOntology ont = manager.loadOntologyFromOntologyDocument(new File("gene_ontology_ext.obo"))
def ontology = ont

def id2class = [:] // maps an OBO-ID to an OWLClass
ontology.getClassesInSignature(true).each {
  def a = it.toString()
  a = a.substring(a.indexOf("obo/")+4,a.length()-1)
  a = a.substring(a.indexOf('#')+1)
  a = a.replaceAll("_",":")
  a = a.replaceAll("tp://bio2rdf.org/","")
  if (id2class[a] == null) {
    id2class[a] = it
  }
}

OWLReasonerFactory reasonerFactory = null

ConsoleProgressMonitor progressMonitor = new ConsoleProgressMonitor()
OWLReasonerConfiguration config = new SimpleConfiguration(progressMonitor)

OWLReasonerFactory fac1 = new ElkReasonerFactory()
OWLReasoner reasoner = fac1.createReasoner(ont,config)

reasoner.precomputeInferences(InferenceType.CLASS_HIERARCHY)

def processes = new TreeSet(reasoner.getSubClasses(id2class["GO:0008150"], false).getFlattened())
def functions = new TreeSet(reasoner.getSubClasses(id2class["GO:0003674"], false).getFlattened())
def cellparts = new TreeSet(reasoner.getSubClasses(id2class["GO:0005575"], false).getFlattened())

def icmap = new LinkedHashMap() // information content map

def double simGIC(Set v1, Set v2, Map ics) { // v1 and v2 are sets of indices, ics is info content map
  def inter = 0.0
  def un = 0.0
  v1.each { 
    if (v2.contains(it)) {
      inter+= (ics[it])
    }
    un+= (ics[it])
  }
  v2.each { un+= (ics[it]) }
  un-=inter
  if (un == 0.0) {
    return 0.0
  } else {
    return inter/un
  }
}

def double simGIC(Set v1, Set v2) { // v1 and v2 are sets of indices
  def inter = 0.0
  def un = 0.0
  v1.each { 
    if (v2.contains(it)) {
      inter+= 1
    }
    un+= 1
  }
  v2.each { un+= 1 }
  un-=inter
  if (un == 0.0) {
    return 0.0
  } else {
    return inter/un
  }
}



def map1 = [:]
def map2 = [:]

def newmap1 = [:]
def newmap2 = [:]

def f1 = new File(opt.g)
f1.splitEachLine("\t") { line ->
  def id = line[1]
  def go = line[4]
  def ev = line[6]
  def tr = line[8]
  if (ev == "EXP" || ev == "IDA" || ev == "IPI" || ev == "IMP" || ev == "IGI" || ev == "IEP") {
    if (tr == "P" || tr == "F") {
    if (go!=null) {
      if (map1[id] == null) {
	map1[id] = new TreeSet()
      }
      map1[id].add(go)
      if (newmap1[id] == null) {
	newmap1[id] = new TreeSet()
      }
      newmap1[id].add(go)
          }
    }
  }
}
def f2 = new File(opt.g)
f2.splitEachLine("\t") { line ->
  def id = line[1]
  def go = line[4]
  def ev = line[6]
  def tr = line[8]
  if (ev == "EXP" || ev == "IDA" || ev == "IPI" || ev == "IMP" || ev == "IGI" || ev == "IEP") {
    //    if (tr == "P") {
    if (tr == "P" || tr == "F") {
    if (go!=null) {
      if (map2[id] == null) {
	map2[id] = new TreeSet()
      }
      map2[id].add(go)
      if (newmap2[id] == null) {
	newmap2[id] = new TreeSet()
      }
      newmap2[id].add(go)
    }
    }
  }
}
def newset = new TreeSet() // set of ids that have new GO annotations
f1 = new File(opt.n)
f1.splitEachLine("\t") { line ->
  def id = line[0]
  def go = line[1]
  if (newmap1[id] == null) {
    newmap1[id] = new TreeSet()
  }
  newmap1[id].add(go)
  newset.add(id)
}
f2 = new File(opt.n)
f2.splitEachLine("\t") { line ->
  def id = line[0]
  def go = line[1]
  if (newmap2[id] == null) {
    newmap2[id] = new TreeSet()
  }
  newmap2[id].add(go)
  newset.add(id)
}

map1.keySet().each { k1 ->
  def v1 = map1[k1]
  def n = new TreeSet()
  v1 = v1.each { go -> 
    if (id2class[go]!=null) {
      reasoner.getSuperClasses(id2class[go], false).getFlattened().each { n.add(it) }
    }
  }
  map1[k1] = n
}

map2.keySet().each { k1 ->
  def v1 = map2[k1]
  def n = new TreeSet()
  v1 = v1.each { go -> 
    if (id2class[go]!=null) {
      reasoner.getSuperClasses(id2class[go], false).getFlattened().each { n.add(it) }
    }
  }
  map2[k1] = n
}

newmap1.keySet().each { k1 ->
  def v1 = newmap1[k1]
  def n = new TreeSet()
  v1 = v1.each { go -> 
    if (id2class[go]!=null) {
      reasoner.getSuperClasses(id2class[go], false).getFlattened().each { n.add(it) }
    }
  }
  newmap1[k1] = n
}

newmap2.keySet().each { k1 ->
  def v1 = newmap2[k1]
  def n = new TreeSet()
  v1 = v1.each { go -> 
    if (id2class[go]!=null) {
      reasoner.getSuperClasses(id2class[go], false).getFlattened().each { n.add(it) }
    }
  }
  newmap2[k1] = n
}


println "Calculating information content"
def counting = [:]
newmap1.keySet().each { key ->
  def val = newmap1[key]
  val.each {
    if (counting[it]==null) {
      counting[it] = 1
    } else {
      counting[it] += 1
    }
  }
}
newmap2.keySet().each { key ->
  def val = newmap2[key]
  val.each {
    if (counting[it]==null) {
      counting[it] = 1
    } else {
      counting[it] += 1
    }
  }
}
def total = newmap1.size()+newmap2.size()

counting.keySet().each { key ->
  def value = counting[key]
  def ic = -Math.log(value/total)/Math.log(2)
  icmap[key] = ic
}

def icmax = 0.0
icmap.keySet().each { key ->
  if (icmap[key]>icmax) {
    icmax = icmap[key]
  }
}

def filterList(Set s1, Set s2) { // filters for processes, functions, cellparts; returns filtered list
  s1.intersect(s2)
}

/* println "Filter maps for processes"
def bpmap1 = map1.collectEntries { key, value -> [key: value.intersect(processes)] }
def bpmap2 = map2.collectEntries { key, value -> [key: value.intersect(processes)] }
def bpnewmap1 = newmap1.collectEntries { key, value -> [key: value.intersect(processes)] }
def bpnewmap2 = newmap2.collectEntries { key, value -> [key: value.intersect(processes)] } */

println "Calculating similarity values"
def pout = new PrintWriter(new BufferedWriter(new FileWriter(opt.o)))
map1.keySet().each { k1 ->
  if (orthomap[k1] != null) {
    def v1 = map1[k1]
    def nv1 = newmap1[k1]
    def l = []
    if (k1 in newset) {
      map2.keySet().each { k2 ->
	//	if (k2 in newset) {
	def v2 = map2[k2]
	def nv2 = newmap2[k2]
	Expando exp = new Expando()
	exp.id = k2
	
	// 	  def valnormal = defer { simGIC(v1, v2, icmap, icmax) }
	// 	  def valboth = defer { simGIC(nv1, nv2, icmap, icmax) }
	def valnormal = defer { simGIC(v1, v2, icmap) }
	def valboth = defer { simGIC(nv1, nv2, icmap) }
	
	exp.valnormal = valnormal.get()
	exp.valboth = valboth.get()
	
	l << exp
      }
      def l1 = l.sort { it.valnormal }.reverse()
      def l4 = l.sort { it.valboth }.reverse()
      def tot = l.size()
      
      for (int i = 0 ; i < l.size() ; i++ ) {
	def exp1 = l1[i]
	if (exp1.id in orthomap[k1]) {
	  def count = 0
	  while (count < tot && l4[count].id!=exp1.id) {
	    count+=1
	  }
	  def res1 = i/tot
	  def res2 = count/tot
	  
	  pout.println("$k1\t"+exp1.id+"\t$res1\t$res2")
	  println("$k1\t"+exp1.id+"\t$res1\t$res2")
	}
      }
    }
  }
}
pout.flush()
pout.close()

/*println "Map 1: "+map1.size()
println "Map 2: "+map2.size()
println "NewMap 1: "+newmap1.size()
println "NewMap 1: "+newmap2.size()*/

System.exit(0)
